#!/usr/bin/env python
# encoding:utf-8
# @Time   : 2019/11/29
# @Author : 胡茂海
# @Site   : 
# @File   : feature_expand.py

import pandas as pd
import os
import pathlib
import sys

root_path = str(pathlib.Path(os.path.abspath(__file__)).parent.parent)
sys.path.append(root_path)
from utils.util import bayesian_blocks, get_bins
from .feature_utils import expand_month_mean_feature
import numpy as np
from datetime import datetime


class CutBins(object):
    def __init__(self):
        self.bins = []
        self.labels = []

    def fit(self, x, y, bins=None, labels=None):
        """
        :param x:
        :param bins: 分割区间[0,10,18,35,55,120] 不指定，默认等频5等分
        :param labels: 区间名称['幼少','青春','成年','中年','老年']
        :return:
        """
        if bins:
            self.bins = bins
        else:
            self.bins = get_bins(y, x)
        if labels:
            self.labels = labels
        else:
            self.labels = range(len(self.bins) - 1)
        return self

    def transform(self, x):
        if not self.bins and not self.labels:
            raise AttributeError("instance is not fitted yet. Call 'fit' first")
        x = pd.cut(x, self.bins, labels=self.labels)
        return x


class ExpandMonths(object):
    def __init__(self, columns=[[]]):
        assert np.array(columns).ndim == 2, '列名必须是二维列表,[[m01_a,m02_a..]]'
        self.columns = columns

    def fit(self, df):
        assert isinstance(df, pd.DataFrame), 'need a dataframe type to expand'
        self.newdata = expand_month_mean_feature(df, self.columns)
        return self

    def transform(self, x):
        return self.newdata


def get_month(cur_month=None):
    if cur_month is None:
        cur_month = datetime.now().month
    month_list = []
    for i in range(1, 13):
        month_list.append(cur_month)
        if cur_month > 0:
            cur_month -= 1
        if cur_month == 0:
            cur_month = 12
    month_list.reverse()
    return month_list


class TimeFeatEncoder(object):
    def __init__(self, time_columns, drop=False, current_month=None, data_current_month=None):
        """
        :param time_columns: should be agree with df order
        :param drop:
        :param current_month: for caculate
        :param data_current_month:
        """
        self.time_columns = time_columns
        self.drop = drop
        self.current_month = current_month
        self.data_current_month = data_current_month
        self.month_agg_columns_ = None

    def fit(self, X, y=None):
        return self

    def transform(self, df):
        col_index_dict = {}
        for i, v in enumerate(self.time_columns):
            col_index_dict[v] = i

        month_list = get_month(self.current_month)
        data_month_list = get_month(self.data_current_month)
        data_month_index = data_month_list.index(month_list[-1])
        data_month = data_month_list[:(data_month_index + 1)]

        org_column = self.time_columns
        col_names = set([i.split('_')[1] for i in org_column])
        columns = []
        for i in col_names:
            temp_col = []
            for j in org_column:
                if i in j:
                    temp_col.append(j)
            columns.append(temp_col)

        c_names = []
        newdf = np.zeros(shape=(len(df), len(col_names) * 3))
        for col_index, column in enumerate(columns):
            month_data = np.zeros(shape=(len(df), 12))
            column_name = '_'.join(column[0].split('_')[1:])

            for index, each_m in enumerate(month_list):
                for c in column:
                    month = int(c.split('_')[0][1:3])
                    cur_index = col_index_dict.get(c)
                    if each_m == month and each_m in data_month:
                        month_data[:, index] = df.iloc[:, cur_index]

            for res_index, mon_value in enumerate([3,6,12]):
                result = np.mean(month_data[:, -mon_value:], axis=1)
                c_name = f'{mon_value}m_{column_name}_mean'
                c_names.append(c_name)
                newdf[:, col_index * 3 + res_index] = result


        if self.drop == True:
            newdf[np.isnan(newdf)] = 0
            self.month_agg_columns_ = c_names
            return newdf

        df = np.concatenate((df.values, newdf), axis=1)
        df[np.isnan(df)] = 0
        self.month_agg_columns_=self.time_columns+c_names
        return df

    def get_feature_names(self, input_features=None):
        return self.month_agg_columns_
